import json
import os
import random
import requests
from pathlib import Path
from urllib.parse import urlparse
from tqdm import tqdm

def download_image(url, save_path, pbar=None):
    """下载图片到指定路径"""
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        
        with open(save_path, 'wb') as f:
            for chunk in response.iter_content(1024):
                f.write(chunk)
        
        if pbar:
            pbar.set_postfix({'status': '下载成功', 'file': os.path.basename(save_path)})
        return True
    except Exception as e:
        if pbar:
            pbar.set_postfix({'status': '下载失败', 'error': str(e)[:20] + '...'})
        return False

def process_data(input_file, output_file, image_root_dir):
    """处理JSON数据并下载图片"""
    # 创建子目录结构
    os.makedirs(os.path.join(image_root_dir, "train"), exist_ok=True)
    os.makedirs(os.path.join(image_root_dir, "val"), exist_ok=True)
    os.makedirs(os.path.join(image_root_dir, "test"), exist_ok=True)
    
    print("正在读取JSON数据...")
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # 打乱数据顺序
    random.shuffle(data)
    
    # 计算划分点
    total = len(data)
    train_end = int(total * 0.8)
    val_end = train_end + int(total * 0.1)
    
    processed_data = []
    success_count = 0
    fail_count = 0
    
    print("\n开始下载图片并处理数据:")
    with tqdm(data, desc="处理进度", unit="item") as pbar:
        for i, item in enumerate(pbar):
            # 确定数据集类型和保存路径
            if i < train_end:
                dataset_type = "train"
            elif i < val_end:
                dataset_type = "val"
            else:
                dataset_type = "test"
            
            # 从URL获取文件扩展名
            ext = os.path.splitext(urlparse(item["small_pic"]).path)[1] or '.jpg'
            filename = f"{item['id']}{ext}"
            relative_path = f"{dataset_type}/{filename}"  # 相对路径
            absolute_path = os.path.join(image_root_dir, relative_path)  # 绝对路径
            
            # 下载图片
            pbar.set_postfix({'status': '下载中...'})
            success = download_image(item["small_pic"], absolute_path, pbar)
            
            if success:
                # 添加新字段
                new_item = item.copy()
                new_item["image"] = relative_path  # 存储相对路径
                new_item["dataset"] = dataset_type
                processed_data.append(new_item)
                success_count += 1
            else:
                fail_count += 1
    
    # 保存处理后的数据
    print("\n正在保存处理结果...")
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(processed_data, f, ensure_ascii=False, indent=2)
    
    print("\n处理完成！")
    print(f"总条目数: {len(data)}")
    print(f"成功处理: {success_count}")
    print(f"失败处理: {fail_count}")
    print(f"训练集: {train_end}, 验证集: {val_end-train_end}, 测试集: {total-val_end}")
    print(f"结果已保存到: {output_file}")
    print(f"图片目录结构:")
    print(f"  {image_root_dir}/")
    print(f"  ├── train/")
    print(f"  ├── val/")
    print(f"  └── test/")

if __name__ == "__main__":
    # 配置路径
    input_json = "merged_output.json"  # 输入的合并JSON文件
    output_json = "final_dataset.json"  # 处理后的输出JSON文件
    image_root_directory = "dataset_images"  # 图片根目录（将自动创建train/val/test子目录）
    
    # 检查并安装tqdm库
    try:
        from tqdm import tqdm
    except ImportError:
        print("正在安装tqdm库...")
        import subprocess
        subprocess.check_call(["pip", "install", "tqdm"])
        from tqdm import tqdm
    
    process_data(input_json, output_json, image_root_directory)